//! # 限流中间件
//!
//! 提供Axum中间件集成的限流功能

use super::{RateLimiter, RequestContext, RateLimitDecision};
use axum::{
    extract::{Request, State},
    http::{HeaderMap, StatusCode},
    middleware::Next,
    response::Response,
};
use std::sync::Arc;
use tracing::{debug, warn};

/// 限流中间件
pub struct RateLimitMiddleware {
    limiter: Arc<RateLimiter>,
}

impl RateLimitMiddleware {
    /// 创建新的限流中间件
    pub fn new(limiter: Arc<RateLimiter>) -> Self {
        Self { limiter }
    }

    /// 中间件处理函数
    pub async fn handle(
        &self,
        request: Request,
        next: Next,
    ) -> Result<Response, StatusCode> {
        // 提取请求信息
        let method = request.method().to_string();
        let uri = request.uri().path().to_string();
        let headers = request.headers().clone();
        
        // 构建请求上下文
        let context = RequestContext::new(uri, method)
            .with_client_ip(extract_client_ip(&headers));

        // 检查限流
        match self.limiter.is_allowed(&context).await {
            Ok(decision) => {
                if decision.allowed {
                    debug!("Request allowed, remaining: {}", decision.remaining);
                    // 添加限流相关的响应头
                    let mut response = next.run(request).await;
                    add_rate_limit_headers(&mut response, &decision);
                    Ok(response)
                } else {
                    warn!("Request rate limited: {:?}", decision.reason);
                    Err(StatusCode::TOO_MANY_REQUESTS)
                }
            }
            Err(e) => {
                warn!("Rate limit check failed: {}", e);
                // 发生错误时允许请求通过
                Ok(next.run(request).await)
            }
        }
    }
}

/// 提取客户端IP地址
fn extract_client_ip(headers: &HeaderMap) -> std::net::IpAddr {
    // 尝试从X-Forwarded-For获取
    if let Some(forwarded) = headers.get("x-forwarded-for") {
        if let Ok(forwarded_str) = forwarded.to_str() {
            if let Some(first_ip) = forwarded_str.split(',').next() {
                if let Ok(ip) = first_ip.trim().parse() {
                    return ip;
                }
            }
        }
    }

    // 尝试从X-Real-IP获取
    if let Some(real_ip) = headers.get("x-real-ip") {
        if let Ok(ip_str) = real_ip.to_str() {
            if let Ok(ip) = ip_str.parse() {
                return ip;
            }
        }
    }

    // 默认返回localhost
    std::net::IpAddr::V4(std::net::Ipv4Addr::new(127, 0, 0, 1))
}

/// 添加限流相关的响应头
fn add_rate_limit_headers(response: &mut Response, decision: &RateLimitDecision) {
    let headers = response.headers_mut();
    
    // 添加剩余配额
    headers.insert("X-RateLimit-Remaining", decision.remaining.to_string().parse().unwrap());
    
    // 添加重置时间
    if let Ok(reset_timestamp) = decision.reset_time.duration_since(std::time::UNIX_EPOCH) {
        headers.insert("X-RateLimit-Reset", reset_timestamp.as_secs().to_string().parse().unwrap());
    }
}